import json
from tqdm import tqdm

from openai import OpenAI
import os
# 升级方舟 SDK 到最新版本 pip install -U 'volcengine-python-sdk[ark]'
from volcenginesdkarkruntime import Ark
import time
import ast
# import pandas as pd

client = OpenAI(api_key="sk-ba3e09bf944e45feacfa5bc4e91af017", base_url="https://api.deepseek.com")

def get_ds_r1_res(send_mess):
    data = {}
    # new_df = pd.DataFrame(columns=['POI名称', 'POI类别', 'POI描述', 'POI卖点', '年龄范围', '性别', '职业类型', '收入水平', '消费能力等级','婚姻状态', '家庭成员', '宗教偏好','旅游类型', '人群类型','景观偏好', '玩法偏好', '行程安排偏好', '餐饮偏好', '住宿偏好', '用户问题', '正负样本'])
    for key, value in tqdm(send_mess.items()):
        str1 = f"""作为推荐系统专家，请根据POI侧的图片描述，生成POI的图片描述信息，保存POI的特征和一些有用于POI推荐的关键信息，以英文回答，总数不超过75词：
                        POI信息：
                            评论：{send_mess[key]}
                        """

        messages = [
            {
                "role": "system",
                "content": "作为推荐系统专家，请根据POI侧的图片描述，生成POI的图片描述信息，保存POI的特征和一些有用于POI推荐的关键信息，以英文回答，总数不超过75词"
            },
            {
                "role": "user",
                "content": str1
            }
        ]
        response = client.chat.completions.create(
            model="deepseek-chat",
            messages=messages,
            stream=False
        )
        try:
            contents = response.choices[0].message.content
            # print(contents)
            data[key] = contents
            # data[key]['poi_image'] = send_mess[key]
        except:
            pass
        # new_df.to_csv('gen_sample_new.csv', index=False, encoding='utf-8')
        # data[key] = send_mess[key]
        with open('image_compress_TKY.json', 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=4)
    return data

if __name__ == '__main__':
    # 文档：https://www.volcengine.com/docs/82379/1449737
    # send_mess = [
    #     {'poi_name':'五彩滩', 'des':'位于涠洲岛东海岸，是看日出的首选。有海蚀崖、海蚀洞、海蚀平台“三位一体”的奇特海蚀地貌景观。滩上是海蚀崖岩，岩下是海蚀平台，崖上随处可见的海蚀洞。退潮后，层叠起伏的海蚀平台上有青苔、海藻、水洼，以及平台之间的潮池，在太阳的照映下五彩斑斓，特别靓眼。'},
    #     {'poi_name':'春夏·静也海景餐厅', 'des':'春夏·静也海景餐厅，一家坐落于风光旖旎的南湾街，临近多个热门景点如南湾海洋运动公园、石螺口、五彩滩和焰舞剧场的优雅用餐之地。这里不仅提供麻辣、清淡、鲜美等多种口味的佳肴，更以其浪漫、温馨、优雅的环境，成为情侣约会、家庭聚餐、生日庆祝或纪念日的理想选择。餐厅内部装饰融合了复古与艺术的元素，营造出一种时尚而休闲的氛围。无论是亲子用餐还是朋友小聚，春夏·静也都能提供安静舒适的空间，让每一位顾客都能在这里享受美食与宁静的双重愉悦。特别值得一提的是，春夏·静也提供的休闲下午茶，让顾客在品尝美食的同时，也能享受到一段悠闲自在的时光。无论是团餐团建还是同学聚会，这里都能满足不同场合的需求，让聚会更加难忘。春夏·静也海景餐厅，一个将美食与美景完美融合的地方，让每一位到访的客人都能在这里找到属于自己的一片宁静与美好。'}
    # ]

    with open('image_txt_TKY.json', 'r', encoding='utf-8') as file:
        send_mess = json.load(file)
    
    send_mess = get_ds_r1_res(send_mess)

    with open('image_txt_compress_TKY.json', 'w', encoding='utf-8') as f:
        json.dump(send_mess, f, ensure_ascii=False, indent=4)